% By Martin Andreasen, November 2024
% This scripts approximates the solution to the NK model 
% using projection to second order
addpath(genpath(pwd))

%% Cleaning
%clc
close all
clear 
format short
%% Settings
appOrder         = 2;          % max order of the polynomial used for approximation
appOrderVar      = [appOrder;appOrder;appOrder]; % Used appOrder per variable: c_t, pi_t, evf_t
endoMeanGrid     = 0;          % 1 for endogenously recenter the grid given the projection approx.
                               % 0 for using the mean in Perturbation to re-center the grid 
muGrid           = 2;          % mu in Smolyak grid
appOrderPer      = 3;          % The order of the perturbation approximation
gridFactor_x     = 1.5;
numIntegralVer   = 1;          % 1 for using N-dimensional monomial (non-product) integration 
                               % 2 for using Gauss-Hermite integration (3 points per dimension)
accuracyGrid     = 0;          % 1 for accuracy on grid, 0 for accuracy on simulated sample
paramsVersion    = 2;          % 1 for standard EZ, 
                               % 2 for model with constant, 
                               % 3 for model with ALFA = 0, 
                               % 4 for reduced model
                               % 5 for model with constant solved under perfect foresight.
lambdaSMM        = 1           %Weight in objective function to Shrinkage of based on Method of Moments (SMM)
SMMwithCSreg     = 1           %2 for incl. CS and risk-adj CS moments, 
                               %1 for incl. CS moments,
                               %0 for not including CS moments
inclSurveys      = 0;
numCPUs          = 2;

% For the optimizer
optimizer        = 1;          % 1 LM, 2 for matlab LM optimizer, 3 for Nelder Mead, 4 for fminunc
dispOn           = 1           % 1 for displaying intermediate results for the optimizer, else 0
maxNumEval       = 10000;      % the maximum number of functions evals during the optimization
MaxIter          = 100;
tolF             = 1D-7;
tolX             = 1D-7;

%% Model parameters: Baseline case
load([pwd,'/ModelSpecification/setupModel'],'setupModel');
if paramsVersion    == 1
    % Solution with standard EZ preferences
    res    = load(['Model_endYear2019_lambdaSMM',num2str(lambdaSMM),'_SMMwithCSreg',num2str(SMMwithCSreg),'_inclSurveys',num2str(inclSurveys),'_withSE.mat']);
elseif paramsVersion == 2
    % solution with constant in EZ preferences
    res    = load(['Model_u_endYear2019_lambdaSMM',num2str(lambdaSMM),'_SMMwithCSreg',num2str(SMMwithCSreg),'_inclSurveys',num2str(inclSurveys),'_withSE.mat']);
elseif paramsVersion == 3
    % solution without EZ preferences
    res = load(['Model_ALFA0_endYear2019_lambdaSMM',num2str(lambdaSMM),'_SMMwithCSreg',num2str(SMMwithCSreg),'_inclSurveys',num2str(inclSurveys),'.mat']);
elseif paramsVersion == 4
    res = load('Model_u_endYear2019_lambdaSMM1_SMMwithCSreg1_inclSurveys0_SIGMAd0.mat');
elseif paramsVersion == 5
    res = load('Model_u_endYear2019_lambdaSMM1_SMMwithCSreg1_inclSurveys0_PFsolv.mat');
end
params = res.outEstim.model.params;

%% Solving the model by perturbation 
MEXon = 1;
[modelPer,errorMes] = perturbationDSGEmodel(params,appOrderPer,setupModel,MEXon);

%% Constructing the grid. 
% Version 1: from a log-linear approximation
varX     = reshape((eye(modelPer.nx^2)-kron(modelPer.hx,modelPer.hx))\reshape((modelPer.eta*modelPer.eta'),modelPer.nx^2,1),modelPer.nx,modelPer.nx);
diagStdX = max(sqrt(diag(varX)),0.001);
meanXPer = (eye(modelPer.nx)-modelPer.hx)\(0.5*modelPer.hss+0.5*reshape(modelPer.hxx,modelPer.nx,modelPer.nx^2)*varX(:));

Smolyak_elem_iso = Smolyak_Elem_Isotrop(modelPer.nx,muGrid);
Smol_grid_iso    = Smolyak_Grid(modelPer.nx,muGrid,Smolyak_elem_iso); 
xGrid            = Smol_grid_iso.*repmat((gridFactor_x*diagStdX'),size(Smol_grid_iso,1),1);
if endoMeanGrid == 0
    xGrid = xGrid + repmat(meanXPer',size(Smol_grid_iso,1),1);
end

%% Nodes for numerical integration
if numIntegralVer == 1
    %covar = modelPer.eta'*modelPer.eta:
    covar = eye(modelPer.ne);  %standard normal shocks
    [n_nodes,epsi_nodes,weight_nodes] = Monomials_2(modelPer.ne,covar);
    %test
    %[sum(exp(epsi_nodes(:,1)).*weight_nodes) exp(0.5*1)]
elseif numIntegralVer == 2
    GH_points = 5;  %Number of integration points
    %etaTilde = modelPer.eta(modelPer.nx1+1:end,:);
    diagEtaTilde = diag(ones(modelPer.ne));   %standard normal shocks
    [n_nodes,epsi_nodes,weight_nodes] = GaussHermiteNdim(GH_points,modelPer.ne,diagEtaTilde);
    %test: 
    %[sum(exp(epsi_nodes(:,1)).*weight_nodes) exp(0.5*1)]
end
if paramsVersion == 5
    % to get perfect foresight solution
    n_nodes = 1;
    epsi_nodes = zeros(1,modelPer.ne);
    weight_nodes = 1;
end
%% Making the setupProj file and setting initial values for the optimizer
constructSetupProj;
%% The starting values for the optimization
fileStart = ['startValues_appOrder',num2str(appOrder)];
[theta0,vecTheta0,setupProj] = loadThetaStart(fileStart,setupProj);
setupProj.nDim1Theta= size(theta0,1);
setupProj.nDim2Theta= size(theta0,2);

tic
setupProj.objectFuncNum = 1; %for main model
[Q0,resEuler0] = NLSobjectFunc(vecTheta0,setupProj);
Q0'*Q0
toc

%% The projection
tic
setupProj.objectFuncNum = 1; %main model
model = projectionStep(vecTheta0,setupProj);
if endoMeanGrid == 1
    setup.xGrid = model.xGrid; %update of grid
end
setupProj.objectFuncNum = 2; %for bond prices
setupProj.maxMat = 40;  %up to 10 years
%modelYieldNLS = projectionStepBondPrices(model,setupProj);
modelYield = projectionStepBondPricesClosedForm(model,setupProj);
toc

%% Evaluating accuracy
% On a grid
setupAccuracy = setupProj;
if accuracyGrid == 1
    gridFactor_xAccuracy     = 6;
    muGridAccuracy           = modelPer.nx;
    Smolyak_elem_isoAccuracy = Smolyak_Elem_Isotrop(modelPer.nx,muGridAccuracy);
    Smol_grid_isoAccuracy    = Smolyak_Grid(modelPer.nx,muGridAccuracy,Smolyak_elem_isoAccuracy);
    xGridAccuracy            = Smol_grid_isoAccuracy.*repmat((gridFactor_xAccuracy*diagStdX'),size(Smol_grid_isoAccuracy,1),1);
elseif accuracyGrid == 0
    % On a simulated sample path
    numSim = 10000;
    rng(1,'twister')
    shocks = randn(model.ne,numSim);
    model.pruningOn = 0;
    % Simulate using perturbation
    [yt,xt,xPruning]=simul(zeros(modelPer.nx,1),shocks,modelPer.g0,...
        modelPer.h0,modelPer.eta,modelPer.derivs,appOrderPer*0+2,1);
    outSim.x_cu = xt-modelPer.h0;

    % Simulate using projection
    %[outSim]        = simProjection(model,shocks);
    
    xGridAccuracy   = outSim.x_cu;
end
setupAccuracy.xGrid      = xGridAccuracy';
setupAccuracy.appOrder   = setupProj.appOrder;
[~,resEulerProj]         = ResidualEquationProj(model.vecTheta,setupAccuracy);
modelYieldAccuracy       = projectionStepBondPrices(modelYield,setupAccuracy);

MAE   = zeros(1,model.ny);
RMSE  = zeros(1,model.ny);
MaxAE = zeros(1,model.ny);
for i=1:model.ny
    MAE(1,i)   = mean(abs(resEulerProj(:,i)));
    RMSE(1,i)  = sqrt(mean(abs(resEulerProj(:,i)).^2));
    MaxAE(1,i) = max(abs(resEulerProj(:,i)));
end

%% Comparing to perturbation
fileStart = ['startValues_appOrder',num2str(appOrder)];
setupPer  = setupAccuracy;
setupPer.appOrder    = appOrderPer;
setupPer.appOrderVar = setupPer.appOrderVar*0+appOrderPer;
[theta0,vecTheta0,setupPer] = loadThetaStart(fileStart,setupPer);
setupPer.nDim1Theta  = size(theta0,1);
setupPer.nDim2Theta  = size(theta0,2);
setupPer.appOrderVar = setupPer.appOrderVar*0+appOrderPer;
setupPer.objectFuncNum = 1; %main model
[~,resEulerPer] = ResidualEquationProj(vecTheta0,setupPer);
MAEPer   = zeros(1,model.ny);
RMSEPer  = zeros(1,model.ny);
MaxAEPer = zeros(1,model.ny);
for i=1:model.ny
    MAEPer(1,i)   = mean(abs(resEulerPer(:,i)));
    RMSEPer(1,i)  = sqrt(mean(abs(resEulerPer(:,i)).^2));
    MaxAEPer(1,i) = max(abs(resEulerPer(:,i)));
end

% Comparing and saving results
results.labelsCol        = [{'MAE'},{'RMSE'},{'MaxAE'}];
results.appOrderVar      = appOrderVar';
results.Proj             = [MAE; RMSE; MaxAE];
results.Per              = [MAEPer; RMSEPer; MaxAEPer];
results.appOrderPer      = appOrderPer;
results.ProjBonds        = [modelYieldAccuracy.MAEbonds modelYieldAccuracy.RMSEbonds modelYieldAccuracy.MaxAEbonds  ];
results.Proj
results.Per
results.Per./results.Proj
mean(log10(results.ProjBonds),1)
%error('stop')
%% Test plotting
%setupProj.modelPer = perturbationDSGEmodel(params,appOrderPer,setupModel,MEXon);
plotProjectionApp(model,setupProj)


%% Simulate the model with projection 
%numSim = 2D5;
%rng(1,'twister')
%shocks = randn(model.ne,numSim);
modelYield.pruningOn = 1;
outSimPrun    = simProjection(modelYield,shocks);
yieldsSimPrun = getYields(modelYield,outSimPrun.xAll_cu);
modelYield.pruningOn = 0;
outSim        = simProjection(modelYield,shocks);
yieldsSim     = getYields(modelYield,outSim.x_cu);
results.meanYPrun      = mean(outSimPrun.y_cu,2);
results.meanY          = mean(outSim.y_cu,2);
results.meanYieldsPrun = mean(yieldsSimPrun,2)*4;
results.meanYields     = mean(yieldsSim,2)*4;
results.stdYPrun       = std(outSimPrun.y_cu,0,2);
results.stdY           = std(outSim.y_cu,0,2);
results.stdYieldsPrun  = std(yieldsSimPrun,0,2)*4;
results.stdYields      = std(yieldsSim,0,2)*4;
results.model          = model;
results.xGridAccuracy  = setupAccuracy.xGrid;

%% Saving the results
if paramsVersion    == 1
    filename = ['AccuracyGrid',num2str(accuracyGrid),'_endYear2019_lambdaSMM',...
        num2str(lambdaSMM),'_SMMwithCSreg',num2str(SMMwithCSreg),'_inclSurveys',num2str(inclSurveys),'_withSE.mat'];
elseif paramsVersion    == 2
    filename = ['AccuracyGrid',num2str(accuracyGrid),'_u_endYear2019_lambdaSMM',...
        num2str(lambdaSMM),'_SMMwithCSreg',num2str(SMMwithCSreg),'_inclSurveys',num2str(inclSurveys),'_withSE.mat'];
elseif paramsVersion    == 3
    filename = ['AccuracyGrid',num2str(accuracyGrid),'_ALFA0_endYear2019_lambdaSMM',...
        num2str(lambdaSMM),'_SMMwithCSreg',num2str(SMMwithCSreg),'_inclSurveys',num2str(inclSurveys),'.mat'];
elseif paramsVersion    == 4
    filename = ['AccuracyGrid_u_SIGMAd0_endYear2019_lambdaSMM',...
        num2str(1),'_SMMwithCSreg',num2str(1),'_inclSurveys',num2str(0),'.mat'];
elseif paramsVersion    == 5
    filename = ['AccuracyGrid_u_PFsolv_endYear2019_lambdaSMM',...
        num2str(1),'_SMMwithCSreg',num2str(1),'_inclSurveys',num2str(0),'.mat'];
end

% Adding extra controls - must appear last due to way the euler errors are computed
modelYield = ProjLoadingsExtraControls(modelYield);
save(filename,'results','modelYield');

